#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
C服务器测试脚本: 负载均衡场景 (场景4.1-4.3)
测试多G服务器的负载均衡和故障切换
"""

import requests
import json
import time
import sys
import os
from typing import Dict, Any, List

class LoadBalanceTester:
    """负载均衡测试类"""
    
    def __init__(self):
        self.gateway_url = "https://www.gongjuxiang.work"
        self.webhook_url = f"{self.gateway_url}/webhook/register"
        self.enhance_url = f"{self.gateway_url}/api/v1/enhance"
        self.correct_secret = "gpu-server-register-to-api-gateway-2024"
        
        # 模拟多个G服务器
        self.g_servers = [
            {"internal_ip": "10.5.0.191", "public_ip": "140.143.155.42", "name": "G1"},
            {"internal_ip": "10.6.0.123", "public_ip": "140.143.155.43", "name": "G2"},
            {"internal_ip": "10.7.0.456", "public_ip": "140.143.155.44", "name": "G3"}
        ]
        
        self.test_image_url = "https://cos.ap-beijing.myqcloud.com/photoenhancei-1259206939/test/small.jpg"
    
    def register_g_server(self, g_server: Dict[str, Any]) -> bool:
        """注册G服务器"""
        register_data = {
            "internal_ip": g_server["internal_ip"],
            "public_ip": g_server["public_ip"],
            "region": "ap-beijing",
            "port": 5800,
            "secret": self.correct_secret
        }
        
        try:
            response = requests.post(self.webhook_url, json=register_data, timeout=10)
            if response.status_code == 200:
                result = response.json()
                print(f"✅ {g_server['name']} 注册成功: {result.get('server_id', '')}")
                return True
            else:
                print(f"❌ {g_server['name']} 注册失败: {response.status_code}")
                return False
        except Exception as e:
            print(f"❌ {g_server['name']} 注册异常: {e}")
            return False
    
    def send_enhance_request(self) -> tuple[bool, Dict[str, Any]]:
        """发送图片处理请求"""
        try:
            data = {"image_url": self.test_image_url}
            response = requests.post(self.enhance_url, json=data, timeout=30)
            
            if response.status_code == 200:
                result = response.json()
                return True, result
            else:
                return False, {"error": response.text, "status_code": response.status_code}
        except Exception as e:
            return False, {"error": str(e)}
    
    def test_scenario_4_1(self):
        """场景4.1: 单G服务器处理所有请求"""
        print("\n" + "="*60)
        print("🧪 场景4.1: 单G服务器处理所有请求")
        print("="*60)
        
        # 只注册一个G服务器
        if not self.register_g_server(self.g_servers[0]):
            return False
        
        time.sleep(2)  # 等待注册完成
        
        # 发送多个请求
        print("📤 发送5个请求...")
        results = []
        for i in range(5):
            success, result = self.send_enhance_request()
            results.append((i+1, success, result))
            print(f"   请求{i+1}: {'✅ 成功' if success else '❌ 失败'}")
            time.sleep(1)
        
        # 验证所有请求都由同一个G服务器处理
        success_count = sum(1 for _, success, _ in results if success)
        
        if success_count >= 3:  # 至少3个成功
            print(f"✅ {success_count}/5 个请求成功")
            print("✅ 场景4.1测试通过")
            return True
        else:
            print(f"❌ 只有{success_count}/5 个请求成功")
            return False
    
    def test_scenario_4_2(self):
        """场景4.2: 多G服务器轮询调度"""
        print("\n" + "="*60)
        print("🧪 场景4.2: 多G服务器轮询调度")
        print("="*60)
        
        # 注册多个G服务器
        print("📋 注册多个G服务器...")
        for g_server in self.g_servers:
            self.register_g_server(g_server)
            time.sleep(1)
        
        time.sleep(3)  # 等待所有注册完成
        
        # 发送多个请求，观察轮询
        print("📤 发送6个请求观察轮询...")
        results = []
        for i in range(6):
            success, result = self.send_enhance_request()
            results.append((i+1, success, result))
            print(f"   请求{i+1}: {'✅ 成功' if success else '❌ 失败'}")
            time.sleep(2)
        
        # 验证请求分发
        success_count = sum(1 for _, success, _ in results if success)
        
        if success_count >= 4:  # 至少4个成功
            print(f"✅ {success_count}/6 个请求成功")
            print("✅ 场景4.2测试通过")
            return True
        else:
            print(f"❌ 只有{success_count}/6 个请求成功")
            return False
    
    def test_scenario_4_3(self):
        """场景4.3: 某个G不健康时自动切换"""
        print("\n" + "="*60)
        print("🧪 场景4.3: 某个G不健康时自动切换")
        print("="*60)
        
        # 注册多个G服务器
        print("📋 注册多个G服务器...")
        for g_server in self.g_servers:
            self.register_g_server(g_server)
            time.sleep(1)
        
        time.sleep(3)  # 等待所有注册完成
        
        # 发送请求验证正常工作
        print("📤 发送请求验证正常工作...")
        success, _ = self.send_enhance_request()
        if not success:
            print("❌ 初始请求失败")
            return False
        
        print("✅ 初始请求成功")
        
        print("⚠️ 注意：请手动停止其中一个G服务器，然后按Enter继续...")
        input()
        
        # 发送请求验证故障切换
        print("📤 发送请求验证故障切换...")
        success_count = 0
        for i in range(3):
            success, result = self.send_enhance_request()
            if success:
                success_count += 1
            print(f"   请求{i+1}: {'✅ 成功' if success else '❌ 失败'}")
            time.sleep(2)
        
        if success_count >= 2:  # 至少2个成功
            print(f"✅ {success_count}/3 个请求成功")
            print("✅ 故障切换正常")
            print("✅ 场景4.3测试通过")
            return True
        else:
            print(f"❌ 只有{success_count}/3 个请求成功")
            return False
    
    def run_all_tests(self):
        """运行所有负载均衡场景测试"""
        print("🚀 开始负载均衡场景测试")
        
        tests = [
            ("场景4.1", self.test_scenario_4_1),
            ("场景4.2", self.test_scenario_4_2),
            ("场景4.3", self.test_scenario_4_3)
        ]
        
        results = []
        for test_name, test_func in tests:
            try:
                print(f"\n🔄 开始执行 {test_name}...")
                result = test_func()
                results.append((test_name, result))
                
                if not result:
                    print(f"⚠️ {test_name} 测试失败，是否继续？(y/n)")
                    if input().lower() != 'y':
                        break
                        
            except Exception as e:
                print(f"❌ {test_name}测试异常: {e}")
                results.append((test_name, False))
        
        # 输出测试结果
        print("\n" + "="*60)
        print("📊 测试结果汇总")
        print("="*60)
        
        passed = 0
        for test_name, result in results:
            status = "✅ 通过" if result else "❌ 失败"
            print(f"{test_name}: {status}")
            if result:
                passed += 1
        
        print(f"\n📈 总体结果: {passed}/{len(results)} 个测试通过")
        
        if passed == len(results):
            print("🎉 所有负载均衡场景测试通过！")
            return True
        else:
            print("⚠️ 部分测试失败，请检查相关配置")
            return False

def main():
    """主函数"""
    print("🧪 C服务器测试脚本: 负载均衡场景")
    print("测试场景: 4.1-4.3")
    print("-" * 60)
    
    tester = LoadBalanceTester()
    success = tester.run_all_tests()
    
    if success:
        print("\n✅ 所有测试完成，结果正常")
        sys.exit(0)
    else:
        print("\n❌ 测试失败，请检查配置")
        sys.exit(1)

if __name__ == "__main__":
    main()
